#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import pymysql
import threading

from datetime import datetime, timedelta
from flask import Flask, request, jsonify
from flask_cors import CORS
from werkzeug.utils import secure_filename

from ragUtil import build_prompt, get_completion
from zhipuUtil import zhipu_util, generate_test_case_by_zhipu
from qwenUtil import qwen_util, generate_test_case_by_qwen
from readPdfAndWriteByMilnus import process_pdf_and_insert_to_milvus
from searchByMilvus import query_milvus_by_text

app = Flask(__name__)

CORS(app, resources={r"/*": {"origins": "*"}})  # 显式指定所有路由前缀为/api的允许所有来源的跨域请求

# 使用IP地址作为键的请求频率字典
ip_request_times = {}

# 请求频率限制参数
MAX_REQUESTS_PER_MINUTE = 5

# 数据库连接配置
db_config = {
    'host': '127.0.0.1',
    'port': 3306,
    'user': 'root',
    'password': '1234qwer',
    'database': 'platform_db',
    'charset': 'utf8mb4'
}


def insert_chat_log(user_id, user_message, assistant_message=None):
    # 连接数据库
    connection = pymysql.connect(**db_config)

    try:
        with connection.cursor() as cursor:
            # 构建SQL语句，注意使用%操作符防止SQL注入
            sql = """
                INSERT INTO user_chat_logs (user_id, user_message, assistant_message, created_at)
                VALUES (%s, %s, %s, %s)
            """
            # 获取当前时间
            current_time = datetime.now()

            # 准备数据，如果assistant_message未提供，默认为NULL
            values = (user_id, user_message, assistant_message if assistant_message is not None else None, current_time)

            # 执行插入操作
            cursor.execute(sql, values)

        # 提交事务
        connection.commit()
    finally:
        # 关闭连接
        connection.close()


def is_request_frequent(ip_address):
    """
    检查基于IP地址的请求是否过于频繁。
    """
    now = datetime.now()
    if ip_address not in ip_request_times:
        # 第一次请求，初始化数据
        ip_request_times[ip_address] = {'count': 1, 'last_request': now}
        return False
    else:
        last_request_time = ip_request_times[ip_address]['last_request']
        time_since_last_request = now - last_request_time
        if time_since_last_request > timedelta(minutes=1):
            # 超过一分钟，重置计数
            ip_request_times[ip_address] = {'count': 1, 'last_request': now}
        else:
            # 更新计数和最后请求时间
            ip_request_times[ip_address]['count'] += 1
            ip_request_times[ip_address]['last_request'] = now

        if ip_request_times[ip_address]['count'] > MAX_REQUESTS_PER_MINUTE:
            return True
        else:
            return False


@app.route('/ask', methods=['POST'])
def ask():
    client_ip = request.remote_addr  # 获取请求者的IP地址

    if is_request_frequent(client_ip):
        return jsonify({'error': '请求过于频繁，请稍后再试'}), 429

    data = request.get_json()
    # 假设用户直接提供了content，没有指定role
    user_content = data.get('content')  # 使用get方法避免KeyError
    model_name = data.get('modelName')
    # insert_chat_log(1, user_content, answer)
    try:
        res = {}
        if model_name == 'qwen':
            res = qwen_util(user_content)
        elif model_name == 'zhipu':
            res = zhipu_util(user_content)
        elif model_name == 'test':
            '''
            根据以下需求文档生成一个全面的测试用例：

                需求文档：
                {user_content}

                测试用例应包括：
                - 功能梳理
                - 前置条件
                - 测试步骤
                - 预期结果

            
                注意：每一个小标题按照用户提供的格式整理
            '''
            '''
            根据以下需求文档生成一个全面且具体的测试用例：

            需求文档：
            {requirement_document}

            测试用例应包括以下部分，并严格按照用户提供的格式整理：

            - 功能梳理：请详细描述功能，包含主要功能点和子功能点。
            - 前置条件：列出所有需要满足的前提条件，包括用户状态、数据准备等。
            - 测试步骤：逐步详细描述测试执行的具体操作步骤，每一步骤应清晰明确。
            - 预期结果：针对每个测试步骤，详细描述预期结果。

            请确保测试用例的每个部分都详尽、具体，便于测试工程师直接使用。
            '''
            prompt = f"""

            根据以下需求文档生成一个全面且具体的测试用例：

            需求文档：
            {user_content}

            测试用例应包括以下部分，并严格按照用户提供的格式整理：

            - 功能梳理：请详细描述功能，包含主要功能点和子功能点。
            - 前置条件：列出所有需要满足的前提条件，包括用户状态、数据准备等。
            - 测试步骤：逐步详细描述测试执行的具体操作步骤，每一步骤应清晰明确。
            - 预期结果：针对每个测试步骤，详细描述预期结果。

            请确保测试用例的每个部分都详尽、具体，便于测试工程师直接使用。
                """
            res = generate_test_case_by_qwen(prompt)
        else:
            res['answer'] = '暂不支持该模型'
        insert_chat_log(client_ip, user_content, res['answer'])
        return jsonify(res)
    except Exception as e:
        print(e)
        answer = '出错了，请稍后再试'
        return jsonify({'answer': answer}), 500


UPLOAD_FOLDER = '/Users/admin/Downloads/uploads'
ALLOWED_EXTENSIONS = {'pdf'}

app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER


def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


@app.route('/upload-pdf', methods=['POST'])
def upload_pdf():
    client_ip = request.remote_addr  # 获取请求者的IP地址

    if is_request_frequent(client_ip):
        return jsonify({'error': '请求过于频繁，请稍后再试'}), 429
    if 'file' not in request.files:
        return 'No file part', 400
    file = request.files['file']
    if file.filename == '':
        return 'No selected file', 400
    if file and allowed_file(file.filename):
        filename = secure_filename(file.filename)
        file.save(os.path.join(app.config['UPLOAD_FOLDER'], filename))
        pdf_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        print(pdf_path)
        try:
            process_pdf_and_insert_to_milvus(pdf_path)
            return 'PDF uploaded and processed successfully.', 200
        except Exception as e:
            return f'Error processing PDF: {str(e)}', 500
    else:
        return 'Invalid file type. Only PDF files are allowed.', 400


@app.route('/query', methods=['POST'])
def query_api():
    try:
        client_ip = request.remote_addr  # 获取请求者的IP地址

        if is_request_frequent(client_ip):
            return jsonify({'error': '请求过于频繁，请稍后再试'}), 429
            # 获取用户查询（假设查询作为JSON的post请求体一部分发送）
        user_query = request.json.get('query', '')

        # 1. 检索 Milvus
        search_results = query_milvus_by_text(user_query)

        # 2. 构建 Prompt
        prompt_template = """
                你是一个问答机器人。
                你的任务是根据下述给定的已知信息回答用户问题。

                已知信息:
                {context}

                用户问：
                {query}

                如果已知信息不包含用户问题的答案，或者已知信息不足以回答用户的问题，请直接回复"我无法回答您的问题"。
                请不要输出已知信息中不包含的信息或答案。
                请用中文回答用户问题。
            """
        prompt = build_prompt(prompt_template, context=search_results, query=user_query)

        # 3. 获取完成的回复
        response = get_completion(prompt)

        # 准备返回的数据
        return_data = {
            "prompt": prompt,
            "answer": response
        }

        return jsonify(return_data), 200

    except Exception as e:
        # 错误处理
        error_message = f"An error occurred: {str(e)}"
        return jsonify({"error": error_message}), 500


#
# @app.route('/test', methods=['POST'])
# def test_api():
#         try:
#             client_ip = request.remote_addr  # 获取请求者的IP地址
#
#             if is_request_frequent(client_ip):
#                 return jsonify({'error': '请求过于频繁，请稍后再试'}), 429
#             # 获取用户查询（假设查询作为JSON的post请求体一部分发送）
#             requirement_document = request.json.get('document', '')
#             prompt = f"""
#                 根据以下需求文档生成一个全面且具体的测试用例：
#
#                 需求文档：
#                 {requirement_document}
#
#                 测试用例应包括以下部分，并严格按照用户提供的格式整理：
#
#                 - 功能梳理：请详细描述功能，包含主要功能点和子功能点。
#                 - 前置条件：列出所有需要满足的前提条件，包括用户状态、数据准备等。
#                 - 测试步骤：逐步详细描述测试执行的具体操作步骤，每一步骤应清晰明确。
#                 - 预期结果：针对每个测试步骤，详细描述预期结果。
#
#                 请确保测试用例的每个部分都详尽、具体，便于测试工程师直接使用。
#
#                 """
#             response=  generate_test_case_by_zhipu(prompt)
#             return_data = {
#                 "prompt": prompt,
#                 "answer": response
#             }
#             return jsonify(return_data), 200
#         except Exception as e:
#             # 错误处理
#             error_message = f"An error occurred: {str(e)}"
#             return jsonify({"error": error_message}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=30001, debug=True)
